import sys
import face_recognition
import os
import multiprocessing
from functools import partial
import json
from typing import List, Dict, Any, Tuple, Optional
import argparse
from PIL import Image
import numpy as np
import logging
import codecs
import time
from multiprocessing import Manager

# 确保Windows控制台使用UTF-8编码
if sys.platform.startswith('win'):
    # 设置控制台输出编码为UTF-8
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer)

# 创建日志处理器
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
# 设置日志格式
formatter = logging.Formatter(
    '%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
console_handler.setFormatter(formatter)

# 配置根日志记录器
logging.basicConfig(
    level=logging.INFO,
    handlers=[console_handler]
)
logger = logging.getLogger(__name__)

class UnknownFaceFound(Exception):
    """当发现未知人脸时抛出的异常，用于中断当前处理流程"""
    pass

class FaceEncodingCache:
    def __init__(self, manager=None):
        self._encodings = []
        self._names = []
        if manager is None:
            self._processed_files = set()
            self._current_batch = set()
            self._saved_faces = set()
            self._face_encodings = []  # 新增：存储已保存人脸的编码
        else:
            self._processed_files = manager.dict()
            self._current_batch = manager.dict()
            self._saved_faces = manager.dict()
            self._face_encodings = manager.list()  # 新增：使用共享列表
        
    def add_face(self, encoding, name, file_path=None):
        """添加新的人脸编码和对应名称"""
        self._encodings.append(encoding)
        self._names.append(name)
        if file_path:
            self._processed_files[file_path] = True
    
    def is_face_saved(self, filename):
        """检查人脸文件是否已保存"""
        return filename in self._saved_faces
    
    def add_saved_face(self, filename):
        """记录已保存的人脸文件"""
        self._saved_faces[filename] = True
    
    def start_new_batch(self):
        """开始新的批次处理"""
        if isinstance(self._current_batch, dict):
            self._current_batch.clear()
        else:
            self._current_batch.clear()
    
    def add_to_current_batch(self, file_path):
        """添加文件到当前批次"""
        self._current_batch[file_path] = True
    
    def is_processed(self, file_path):
        """检查文件是否已经处理过"""
        return file_path in self._processed_files or file_path in self._current_batch

    def get_faces(self):
        """获取所有已缓存的人脸编码和名称"""
        return self._encodings.copy(), self._names.copy()
    
    @property
    def size(self):
        """获取缓存中的人脸数量"""
        return len(self._names)
    
    def clear(self):
        """清空缓存"""
        self._encodings.clear()
        self._names.clear()
        if isinstance(self._processed_files, set):
            self._processed_files.clear()
        else:
            while self._processed_files:
                self._processed_files.pop()

    def add_face_encoding(self, encoding):
        """添加人脸编码到缓存"""
        self._face_encodings.append(encoding)
    
    def find_similar_face(self, encoding, similarity_tolerance):
        """查找是否存在相似的人脸
        
        Args:
            encoding: 要比较人脸编码
            similarity_tolerance: 相似度阈值，越小越严格
            
        Returns:
            bool: 是否找到相似人脸
        """
        if not self._face_encodings:
            return False
            
        # 计算与所有已知人脸的距离
        distances = face_recognition.face_distance(self._face_encodings, encoding)
        return any(d <= similarity_tolerance for d in distances)

class FileInfo:
    """存储文件信息的类
    
    Attributes:
        file_name: 文件名
        file_path: 规范化后的文件完整路径
        people_names: 在图片中识别到的人名列表
    """
    def __init__(self, file_name: str, file_path: str):
        self.file_name = file_name
        self.file_path = normalize_path(file_path)
        self.people_names: List[str] = []

    def to_dict(self) -> Dict[str, Any]:
        """将FileInfo对象转换为可序列化的字典格式"""
        return {
            "fileName": self.file_name,
            "filePath": self.file_path,
            "peopleNames": self.people_names if self.people_names else ["未知"]
        }

def normalize_path(path: str) -> str:
    """统一路径分隔符为 '/'
    
    Args:
        path: 原始路径
        
    Returns:
        统一使用 '/' 作为分隔符的路径
    """
    return path.replace('\\', '/')

def verify_face_encoding(image_path: str) -> bool:
    """验证是否能从图片中提取人脸特征
    
    Args:
        image_path: 图片文件路径
        
    Returns:
        bool: 如果能提取到人脸特征返回True，否则返回False
    """
    try:
        image = face_recognition.load_image_file(image_path)
        face_encodings = face_recognition.face_encodings(image)
        return len(face_encodings) > 0
    except Exception as e:
        logger.error(f"验证人脸特征时发生错误 - 文件: {image_path}, 错误: {str(e)}")
        return False

def save_face_image(
    image: np.ndarray, 
    face_location: tuple, 
    face_encoding: np.ndarray,  # 新增：传入人脸编码
    known_faces_dir: str, 
    original_image_path: str, 
    skipped_files: set,
    face_cache: FaceEncodingCache,
    save_queue: multiprocessing.Queue
) -> Optional[str]:
    """从原图中裁剪人脸并保存为新文件"""
    top, right, bottom, left = face_location
    
    # 计算人脸区域的尺寸和边距（扩展30%）
    face_height = bottom - top
    face_width = right - left
    margin_h = int(face_height * 0.3)
    margin_w = int(face_width * 0.3)
    
    # 确保边距不会超出图片范围
    image_height, image_width = image.shape[:2]
    top = max(0, top - margin_h)
    bottom = min(image_height, bottom + margin_h)
    left = max(0, left - margin_w)
    right = min(image_width, right + margin_w)
    
    # 裁剪并处理人脸图像
    face_image = Image.fromarray(image[top:bottom, left:right])
    
    # 确保图像尺寸足够大（最小224像素）
    min_size = 224
    width, height = face_image.size
    if width < min_size or height < min_size:
        # 如果宽度或高度小于最小尺寸，就需要放大图像
        ratio = max(min_size/width, min_size/height)  # 计算需要放大的比例
        new_size = (int(width*ratio), int(height*ratio))  # 计算新的尺寸
        # 使用LANCZOS算法进行图像重采样放大
        face_image = face_image.resize(new_size, Image.Resampling.LANCZOS)
    
    # 将保存任务添加到队列，同时传入人脸编码
    save_queue.put((face_image, face_encoding, original_image_path))
    return None  # 文件名将由保存进程生成

# 修改load_known_faces函数
def load_known_faces(known_faces_dir: str, cache: FaceEncodingCache) -> Tuple[List[Any], List[str]]:
    """加载已知人脸编码和对应的名称"""
    start_time = time.time()
    
    # 如果缓存为空，则加载所有人脸
    if cache.size == 0:
        logger.info(f"正在从目录加载已知人脸: {known_faces_dir}")
        # 获取目录下所有文件并排序，确保顺序一致
        image_files = sorted(os.listdir(known_faces_dir))
        for image_name in image_files:
            image_path = normalize_path(os.path.join(known_faces_dir, image_name))
            face_load_start = time.time()
            try:
                image = face_recognition.load_image_file(image_path)
                face_encodings = face_recognition.face_encodings(image)
                if face_encodings:
                    cache.add_face(face_encodings[0], os.path.splitext(image_name)[0], image_path)
                    logger.debug(f"已加载人脸: {image_name}, 耗时: {time.time() - face_load_start:.2f}秒")
                else:
                    logger.warning(f"无法从文件中提取人脸特征: {image_name}")
            except Exception as e:
                logger.error(f"加载已人脸时发生错误 - 文件: {image_name}, 错误: {str(e)}")
    
    known_face_encodings, known_face_names = cache.get_faces()
    total_time = time.time() - start_time
    logger.info(f"完成已知人脸加载，共 {len(known_face_names)} 个有效人脸，总耗时: {total_time:.2f}秒")
    return known_face_encodings, known_face_names

def process_image(image_path: str, known_face_encodings: List[Any], 
                 known_face_names: List[str], recognition_tolerance: float, 
                 known_faces_dir: str, skipped_files: set,
                 face_cache: FaceEncodingCache,
                 save_queue: multiprocessing.Queue) -> Tuple[FileInfo, List[Tuple[np.ndarray, str]]]:
    """处理单个图像文件的人脸识别，返回文件信息和新发现的人脸"""
    start_time = time.time()
    new_faces = []  # 存储新发现的人脸编码和文件名
    
    # 添加图片格式检查
    try:
        img = Image.open(image_path)
        # 转换RGBA图片为RGB
        if img.mode == 'RGBA':
            img = img.convert('RGB')
        width, height = img.size
        
        # 如果图片太大，先调整大小
        max_size = 1600  # 设置最大尺寸
        if width > max_size or height > max_size:
            scale = max_size / max(width, height)
            new_size = (int(width * scale), int(height * scale))
            img = img.resize(new_size, Image.Resampling.LANCZOS)
            
        # 转换为numpy数组
        image = np.array(img)
        logger.debug(f"图片已加载: {width}x{height}")
    except Exception as e:
        logger.error(f"处理图片时发生错误: {image_path}, 错误: {str(e)}")
        skipped_files.add(image_path)
        return FileInfo(os.path.basename(image_path), image_path), new_faces

    if image_path in skipped_files:
        logger.debug(f"跳过已知无法处理的图片: {image_path}")
        return FileInfo(os.path.basename(image_path), image_path), new_faces

    file_info = FileInfo(os.path.basename(image_path), image_path)
    face_cache.add_to_current_batch(image_path)
    logger.debug(f"开始处理图片: {image_path}")

    try:
        face_locations = face_recognition.face_locations(image)
        face_encodings = face_recognition.face_encodings(image, face_locations)
    except Exception as e:
        logger.error(f"处理图片时发生错误: {image_path}, 错误: {str(e)}")
        skipped_files.add(image_path)
        return file_info, new_faces

    if not face_encodings:
        logger.debug(f"未在图片中检测到人脸: {image_path}")
        return file_info, new_faces

    for face_encoding, face_location in zip(face_encodings, face_locations):
        matches = face_recognition.compare_faces(known_face_encodings, face_encoding, 
                                               tolerance=recognition_tolerance)
        
        if not any(matches) and known_face_encodings:
            new_face_filename = save_face_image(
                image, face_location, face_encoding,  # 传入face_encoding
                known_faces_dir, image_path, skipped_files, 
                face_cache, save_queue
            )
            if new_face_filename:
                name = os.path.splitext(new_face_filename)[0]
                if name not in file_info.people_names:
                    file_info.people_names.append(name)
                    new_faces.append((face_encoding, name))
                    logger.info(f"成功保存新的未知人脸: {new_face_filename}")
        elif known_face_encodings:
            face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
            best_match_index = face_distances.argmin()
            if matches[best_match_index]:
                name = known_face_names[best_match_index]
                if name not in file_info.people_names:
                    file_info.people_names.append(name)
                    logger.debug(f"识别到已知人脸 {name}")
        else:
            new_face_filename = save_face_image(
                image, face_location, face_encoding,  # 传入face_encoding
                known_faces_dir, image_path, skipped_files,
                face_cache, save_queue
            )
            if new_face_filename:
                name = os.path.splitext(new_face_filename)[0]
                file_info.people_names.append(name)
                new_faces.append((face_encoding, name))
                logger.info(f"成功保存新的未知人脸: {new_face_filename}")

    # logger.info(f"图片 {image_path} 处理完成，总耗时: {time.time() - start_time:.2f}秒")
    return file_info, new_faces

def get_image_files(path: str) -> List[str]:
    """递归获取目录中的所有图像文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif'}
    image_files = []
    
    if os.path.isfile(path):
        if os.path.splitext(path)[1].lower() in image_extensions:
            image_files.append(normalize_path(path))
    elif os.path.isdir(path):
        for root, _, files in os.walk(path):
            for file in files:
                if os.path.splitext(file)[1].lower() in image_extensions:
                    image_files.append(normalize_path(os.path.join(root, file)))
    
    logger.info(f"找到 {len(image_files)} 个图像文件")
    return image_files

def batch_process_images(image_paths: List[str], process_func, face_cache: FaceEncodingCache,
                        known_face_encodings: List[Any], known_face_names: List[str],
                        cpus: int) -> List[FileInfo]:
    """批量处理图片，使用进程池并行处理"""
    results = []
    total_images = len(image_paths)
    processed_count = 0
    
    logger.info(f"开始处理图片，共 {total_images} 个文件")
    
    # 计算每个进程处理的图片数量
    chunk_size = max(1, total_images // (cpus * 4))  # 每个进程处理多个图片以减少进程切换开销
    
    with multiprocessing.Pool(processes=cpus) as pool:
        # 使用imap_unordered提高并行效率
        for file_info, new_faces in pool.imap_unordered(
            process_func, 
            image_paths,
            chunksize=chunk_size
        ):
            processed_count += 1
            results.append(file_info)
            
            # 动态添加新发现的人脸到已知人脸列表
            for encoding, name in new_faces:
                known_face_encodings.append(encoding)
                known_face_names.append(name)
                face_cache.add_face(encoding, name)
            
            # 打印进度
            if processed_count % 10 == 0 or processed_count == total_images:
                logger.info(f"总进度: {processed_count}/{total_images} "
                          f"({processed_count/total_images*100:.1f}%)")
    
    logger.info(f"所有图片处理完成，共处理 {processed_count} 个文件")
    return results

def save_face_worker(save_queue, face_cache, known_faces_dir, duplicate_check_tolerance):
    """专门处理人脸保存的工作进程
    
    Args:
        save_queue: 保存任务队列
        face_cache: 人脸缓存对象
        known_faces_dir: 已知人脸目录
        duplicate_check_tolerance: 用于检查重复人脸的相似度阈值
    """
    saved_faces_count = 0
    skipped_faces_count = 0
    
    # 获取当前最大的未知人脸索引
    existing_files = os.listdir(known_faces_dir)
    face_indices = []
    for filename in existing_files:
        if filename.startswith("未知人脸_") and filename.endswith(".jpg"):
            try:
                index = int(filename[4:-4].split('_')[1])
                face_indices.append(index)
            except (ValueError, IndexError):
                continue
    
    # 设置起始索引为当前最大索引+1
    face_index = max(face_indices, default=0) + 1
    
    while True:
        try:
            # 从队列获取保存任务
            task = save_queue.get()
            if task is None:  # 结束信号
                break
                
            face_image, face_encoding, original_image_path = task
            
            # 首先检查是否存在相似的人脸
            if face_cache.find_similar_face(face_encoding, 
                                          similarity_tolerance=duplicate_check_tolerance):
                logger.info(f"发现相似人脸，跳过保存 - 来自图片: {original_image_path}")
                skipped_faces_count += 1
                continue
            
            # 生成新文件名
            while True:
                new_filename = f"未知人脸_{face_index}.jpg"
                full_path = os.path.join(known_faces_dir, new_filename)
                
                if not os.path.exists(full_path):
                    try:
                        # 保存图片
                        face_image.save(full_path, quality=95)
                        
                        # 验证保存的图片是否可以提取人脸特征
                        if verify_face_encoding(full_path):
                            logger.info(f"成功保存新的未知人脸: {new_filename} "
                                      f"(已保存: {saved_faces_count + 1}, 已跳过: {skipped_faces_count})")
                            face_cache.add_face_encoding(face_encoding)
                            saved_faces_count += 1
                            break
                        else:
                            # 如果无法提取特征，删除文件并记录
                            os.remove(full_path)
                            logger.warning(f"无法从保存的人脸图片中提取特征 - 原始图片: {original_image_path}")
                            skipped_faces_count += 1
                            break
                    except Exception as e:
                        logger.error(f"保存人脸时发生错误: {str(e)}")
                        if os.path.exists(full_path):
                            os.remove(full_path)
                        skipped_faces_count += 1
                        break
                
                face_index += 1
                
        except Exception as e:
            logger.error(f"保存人脸时发生错误: {str(e)}")
            skipped_faces_count += 1
            continue
    
    logger.info(f"人脸保存进程结束，共保存 {saved_faces_count} 个新人脸，"
                f"跳过 {skipped_faces_count} 个重复或无效人脸")

def main():
    total_start_time = time.time()  # 记录程序开始时间
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='人脸识别脚本')
    parser.add_argument('--cpus', type=int, default=1, help='使用的CPU核心数')
    parser.add_argument('--tolerance', type=float, default=0.6, 
                       help='人脸识别匹配的容差值，用于确定是否为同一个人（值越小越严格）')
    parser.add_argument('--show-distance', type=str, choices=['true', 'false'], 
                       default='false', help='显示距离（当前版本未使用）')
    parser.add_argument('known_faces_dir', type=str, help='已知人脸目录')
    parser.add_argument('unknown_faces_input', type=str, help='待识别的图像文件或目录')
    args = parser.parse_args()

    # 初始化参数
    cpus = args.cpus if args.cpus > 0 else multiprocessing.cpu_count()
    known_faces_dir = normalize_path(os.path.expanduser(args.known_faces_dir))
    unknown_faces_input = normalize_path(os.path.expanduser(args.unknown_faces_input))
    
    logger.info(f"使用 {cpus} 个CPU核心进行处理")
    logger.info(f"人脸识别容差设置为 {args.tolerance}")

    # 始化跳过文件集
    skipped_files = set()

    # 使用Manager来创建进程间共享的数据结构
    with Manager() as manager:
        save_queue = manager.Queue()
        face_cache = FaceEncodingCache(manager)

        # 启动保存工作进程，传入known_faces_dir参数
        save_process = multiprocessing.Process(
            target=save_face_worker,
            args=(save_queue, face_cache, known_faces_dir, args.tolerance)
        )
        save_process.start()

        try:
            # 加载已知人脸
            known_face_encodings, known_face_names = load_known_faces(known_faces_dir, face_cache)
            unknown_image_paths = get_image_files(unknown_faces_input)
            
            if not unknown_image_paths:
                logger.error(f"在 '{unknown_faces_input}' 中没有找到有效的图像文件")
                sys.exit(1)

            # 创建处理函数
            process_func = partial(process_image, 
                                known_face_encodings=known_face_encodings,
                                known_face_names=known_face_names,
                                recognition_tolerance=args.tolerance,
                                known_faces_dir=known_faces_dir,
                                skipped_files=skipped_files,
                                face_cache=face_cache,
                                save_queue=save_queue)
            
            # 使用优化后的批处理函数
            file_infos = batch_process_images(
                unknown_image_paths, 
                process_func,
                face_cache,
                known_face_encodings,
                known_face_names,
                cpus
            )

        finally:
            # 发送结束信号并等待保存进程结束
            save_queue.put(None)
            save_process.join()
            
            # 打印总体执行时间
            total_time = time.time() - total_start_time
            hours = int(total_time // 3600)
            minutes = int((total_time % 3600) // 60)
            seconds = total_time % 60
            
            logger.info(f"程序执行完成，总耗时: "
                       f"{hours}小时 {minutes}分钟 {seconds:.1f}秒")

    # 打印被跳过的文件
    if skipped_files:
        logger.warning("以下文件由于无法提取人脸特征被跳过:")
        for file_path in skipped_files:
            logger.warning(f"- {file_path}")

    # 保存结果
    output_file = 'face_recognition_results.json'
    json_results = [file_info.to_dict() for file_info in file_infos]
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(json_results, f, ensure_ascii=False, indent=4)
    
    logger.info(f"人脸识别完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    if sys.platform.startswith('win'):
        multiprocessing.freeze_support()
    main()